﻿using System.Collections;

namespace JESAI.DistributedLock.ZooKeeper;

using JESAI.DistributedLock.Internal;
using org.apache.zookeeper;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

/// <summary>
/// 此模式利用了这样一个事实：在“持有”锁（或其他同步原语）时，连接基本上是空闲的。因此，与其为每个持有的锁创建一个新连接，通常可以复用一个共享连接，以便该连接可以同时持有多个锁。
/// 
/// 复用默认是开启的，除非<see cref="UseTransaction(bool)"/>被设置为TRUE，在这种情况下复用会被禁用，因为它与<see cref="UseTransaction(bool)"/>不兼容。
/// 
/// 这种实现方式使得释放这样一个连接上持有的锁永远不会被等待在同一连接上获取锁的Acquire()调用阻塞。因此，复用策略是“乐观的”：如果在共享连接上不能立即获取锁，将分配一个新的（可共享的）连接。
/// 
/// 此选项可以在高负载场景下提高性能并避免连接池耗尽。它也特别适用于<see cref="IDistributedLock.TryAcquire(TimeSpan, System.Threading.CancellationToken)"/>语义与零长度超时一起使用的情况。
/// </summary>

internal class ZooKeeperConnection : IDisposable
{
    /// <summary>
    /// 希望10分钟的时间能够防止泄漏变得太严重，同时通过允许我们在负载下重用会话来提高效率。
    /// </summary>
    public static readonly Pool DefaultPool = new(maxAge: TimeSpan.FromMinutes(10));

    private readonly InternalConnection _internalConnection;
    private Action? _releaseToPool;

    private ZooKeeperConnection(InternalConnection internalConnection, Action releaseToPool)
    {
        this._internalConnection = internalConnection;
        this._releaseToPool = releaseToPool;
    }

    public ZooKeeper ZooKeeper => this._internalConnection.ZooKeeper;

    public CancellationToken ConnectionLostToken => this._internalConnection.ConnectionLostToken;

    public void Dispose() => Interlocked.Exchange(ref this._releaseToPool, null)?.Invoke();

    public sealed class Pool
    {
        private readonly Dictionary<ZooKeeperConnectionInfo, ConnectionEntry> Connections = new();
        private readonly TimeoutValue _maxAge;

        public Pool(TimeoutValue maxAge)
        {
            Invariant.Require(!maxAge.IsInfinite);
            this._maxAge = maxAge.TimeSpan;
        }

        private object PoolLock => Connections.As<ICollection>().SyncRoot;

        public Task<ZooKeeperConnection> ConnectAsync(ZooKeeperConnectionInfo connectionInfo, CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();

            lock (this.PoolLock)
            {
                // 如果我们有游泳池的入口，就用它
                if (this.Connections.TryGetValue(connectionInfo, out var entry))
                {
                    ++entry.UserCount;
                    return ToResultAsync(entry);
                }

                // 创建新连接
                var newConnectionTask = this.InternalConnectAsync(connectionInfo);
                var newEntry = new ConnectionEntry(newConnectionTask)
                {
                    // 2，因为我们有当前请求和即将创建的超时任务
                    UserCount = 2
                };
                this.Connections.Add(connectionInfo, newEntry);
                newConnectionTask.ContinueWith(OnConnectionTaskCompleted);
                return ToResultAsync(newEntry);

                void OnConnectionTaskCompleted(Task<InternalConnection> internalConnectionTask)
                {
                    // 如果我们从未连接过，就释放我们对任务的控制
                    if (internalConnectionTask.Status != TaskStatus.RanToCompletion)
                    {
                        this.ReleaseEntry(connectionInfo, newEntry, remove: true);
                    }
                    // 否则，请等待最大年龄结束或连接丢失，以释放我们的保留。这确保了连接不会过早地从池中删除，并且即使停止使用，它最终也会被删除
                    else
                    {
                        Task.Delay(this._maxAge.TimeSpan, internalConnectionTask.Result.ConnectionLostToken)
                            .ContinueWith(_ => this.ReleaseEntry(connectionInfo, newEntry, remove: true));
                    }
                }
            }

            async Task<ZooKeeperConnection> ToResultAsync(ConnectionEntry entry)
            {
                try
                {
                    var internalConnection = await entry.ConnectionTask.ConfigureAwait(false);
                    return new ZooKeeperConnection(internalConnection, releaseToPool: () => this.ReleaseEntry(connectionInfo, entry, remove: false));
                }
                catch
                {
                    // 如果我们无法构建连接，仍然可以释放对条目的控制以允许清理
                    this.ReleaseEntry(connectionInfo, entry, remove: false);
                    throw;
                }
            }
        }

        private async Task<InternalConnection> InternalConnectAsync(ZooKeeperConnectionInfo connectionInfo)
        {
            var watcher = new ConnectionWatcher(connectionInfo.SessionTimeout);
            var zooKeeper = new ZooKeeper(
                connectstring: connectionInfo.ConnectionString,
                sessionTimeout: connectionInfo.SessionTimeout.InMilliseconds,
                watcher: watcher
            );

            using var timeoutSource = new CancellationTokenSource(connectionInfo.ConnectTimeout.TimeSpan);
            using var timeoutRegistration = timeoutSource.Token.Register(
                () => watcher.TaskCompletionSource.TrySetException(new TimeoutException($"Timed out connecting to ZooKeeper after {connectionInfo.ConnectTimeout.InMilliseconds}ms"))
            );

            foreach (var authInfo in connectionInfo.AuthInfo)
            {
                zooKeeper.addAuthInfo(authInfo.Scheme, authInfo.Auth.ToArray());
            }

            try
            {
                await watcher.TaskCompletionSource.Task.ConfigureAwait(false);
                return new InternalConnection(zooKeeper, watcher);
            }
            catch
            {
                // 失败时，清理我们创建的实例
                try { await zooKeeper.closeAsync().ConfigureAwait(false); }
                finally { watcher.Dispose(); }
                throw;
            }
        }

        private void ReleaseEntry(ZooKeeperConnectionInfo connectionInfo, ConnectionEntry entry, bool remove)
        {
            bool shouldDispose;
            lock (this.PoolLock)
            {
                if (remove)
                {
                    this.Connections.As<ICollection<KeyValuePair<ZooKeeperConnectionInfo, ConnectionEntry>>>()
                        .Remove(new KeyValuePair<ZooKeeperConnectionInfo, ConnectionEntry>(connectionInfo, entry));
                }

                shouldDispose = --entry.UserCount == 0;
                // 我们将超时监视器任务作为“用户”包含在内，这一事实支持了这一保证
                Invariant.Require(
                    !shouldDispose || !(this.Connections.TryGetValue(connectionInfo, out var registeredEntry) && registeredEntry == entry),
                    "If we're disposing then the entry must be removed from the pool"
                );
            }

            if (shouldDispose)
            {
                //在后台启动连接处理，以避免阻塞/丢弃，并处理任务尚未完成的情况
                entry.ConnectionTask.ContinueWith(
                    t => { var ignored = t.Result.DisposeAsync(); },
                    TaskContinuationOptions.OnlyOnRanToCompletion
                );
            }
        }

        private class ConnectionEntry
        {
            public ConnectionEntry(Task<InternalConnection> connectionTask)
            {
                this.ConnectionTask = connectionTask;
            }

            public Task<InternalConnection> ConnectionTask { get; }
            // 由游泳池的锁保护
            public int UserCount { get; set; }
        }
    }

    private class InternalConnection : IAsyncDisposable
    {
        private readonly ConnectionWatcher _watcher;

        public InternalConnection(ZooKeeper zooKeeper, ConnectionWatcher watcher)
        {
            this.ZooKeeper = zooKeeper;
            this._watcher = watcher;
        }

        public ZooKeeper ZooKeeper { get; }

        public CancellationToken ConnectionLostToken => this._watcher.ConnectionLost;

        public async ValueTask DisposeAsync()
        {
            try { await this.ZooKeeper.closeAsync().ConfigureAwait(false); }
            finally { this._watcher.Dispose(); }
        }
    }

    private class ConnectionWatcher : Watcher, IDisposable
    {
        private readonly CancellationTokenSource _connectionLostSource = new();
        private readonly TimeoutValue _sessionTimeout;
        private int _isWaitingForReconnect;

        public ConnectionWatcher(TimeoutValue sessionTimeout)
        {
            this._sessionTimeout = sessionTimeout;
        }

        public TaskCompletionSource<ValueTuple> TaskCompletionSource { get; } = new TaskCompletionSource<ValueTuple>();
        public CancellationToken ConnectionLost => this._connectionLostSource.Token;

        public override Task process(WatchedEvent @event)
        {
            if (@event.getState() == Event.KeeperState.SyncConnected)
            {
                if (!this.TaskCompletionSource.TrySetResult(default))
                {
                    // 如果我们要重新连接，请清除任何预定的取消
                    this._connectionLostSource.CancelAfter(Timeout.Infinite);
                    Volatile.Write(ref this._isWaitingForReconnect, 0);
                }
            }
            else
            {
                if (!this.TaskCompletionSource.TrySetException(new InvalidOperationException($"Failed to connect to ZooKeeper. State: {@event.getState()}")))
                {
                    // 如果我们看到已过期的事件（如果我们从未重新连接，则可能不会出现），只会断开连接
                    if (@event.getState() == Event.KeeperState.Expired)
                    {
                        this._connectionLostSource.Cancel();
                    }
                    // 否则，动物园管理员将尝试重新连接。在我们开火之前，将其交给会话超时来执行。请参阅https://zookeeper.apache.org/doc/r3.6.2/zookeeperProgrammers.html 
                    else if (Interlocked.Exchange(ref this._isWaitingForReconnect, 1) == 0)
                    {
                        // 只有当我们改变了！重新连接；否则，如果我们收到多个失败事件，我们将继续延长计时器
                        this._connectionLostSource.CancelAfter(this._sessionTimeout.TimeSpan);
                    }
                }
            }

            return Task.CompletedTask;
        }

        public void Dispose() => this._connectionLostSource.Dispose();
    }
}
